package com.wuwangfu.state;

import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.restartstrategy.RestartStrategies;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.calcite.shaded.com.google.common.base.Charsets;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;

import java.io.RandomAccessFile;

/**
 * @Author: jcshen
 * @Date: 2023-03-07
 * @PackageName: com.wuwangfu.state
 * @ClassName: AtLeastOnceSource
 * @Description: OperatorState实现AtLeastOnce
 * @Version: 1.0.0
 *
 */
public class AtLeastOnceSource {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(2);
        env.enableCheckpointing(30000);
        env.setRestartStrategy(RestartStrategies.fixedDelayRestart(3,5000));

        DataStreamSource<String> source = env.addSource(new MyAtLeastOnceSource("F:\\workspace\\flink\\flink-parent\\flink-14\\data"));
        DataStreamSource<String> line = env.socketTextStream("localhost", 8888);
        SingleOutputStreamOperator<String> maped = line.map(new MapFunction<String, String>() {
            @Override
            public String map(String value) throws Exception {
                if ("error".equals(value)) {
                    int v = 10 / 0;
                }
                return value;
            }
        });

        DataStream<String> unioned = source.union(maped);

        unioned.print();

        env.execute();
    }

    private static class MyAtLeastOnceSource extends RichParallelSourceFunction<String>implements CheckpointedFunction {

        private String path;

        public MyAtLeastOnceSource(String path) {
            this.path = path;
        }

        private transient boolean flag = true;
        private transient Long offset = 0L;
        private transient ListState<Long> listState;

        /**
         * @param context
         * @throws Exception
         *
         * 初始化状态或恢复状态执行一次，在 run 方法执行前执行一次
         */
        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            /*保存偏移量*/
            ListStateDescriptor<Long> stateDesc = new ListStateDescriptor<>("offset-state", Long.class);
            listState = context.getOperatorStateStore().getListState(stateDesc);
            //当前的状态是否已恢复
            if (context.isRestored()){
                //从ListState中恢复偏移量
                Iterable<Long> iterable = listState.get();
                for (Long iter : iterable){
                    offset = iter;
                }
            }
        }

        /**
         * @param context
         * @throws Exception
         *
         * 在checkpoint时会执行一次，该方法会周期性的调用
         */
        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            /*定期的更新OperatorState*/
            //清空之前的状态
            listState.clear();
            //更新状态
            listState.add(offset);

        }

        @Override
        public void run(SourceContext<String> ctx) throws Exception {
            int index = getRuntimeContext().getIndexOfThisSubtask();
            RandomAccessFile file = new RandomAccessFile(path + "/" + index + ".txt", "rw");
            //从指定的偏移量读取数据
            file.seek(offset);
            //读取数据
            while (flag){
                String line = file.readLine();
                if (line !=null){
                    line = new String(line.getBytes(Charsets.ISO_8859_1),Charsets.UTF_8);
                    /**
                     * 加锁，锁对象：
                     *      snapshotState中checkpoint时 可能会失败
                     *      确保 run时发送的数据完毕
                     *
                     */
                    synchronized (ctx.getCheckpointLock()){
                        //指针：指向最新的偏移量
                        offset = file.getFilePointer();
                        //输出
                        ctx.collect(index + ":"+ line);
                    }

                }else{
                    Thread.sleep(1000);
                }
            }
        }

        @Override
        public void cancel() {
            flag = false;
        }


    }
}
